"""
Borrowed code from Steven's giza corpus reader, building naacl corpus reader based on that.

brianjing404, bkbk.var
"""

import re
import os


def read_block_giza(stream):
    s = ''
    while True:
        line = stream.readline()
        if line[0] == '=' or line[0] == '\n' or line[:2] == '\r\n':
            continue
        # End of file:
        if not line:
            if s: return [s]
            else: return []
        # Other line:
        else:
            s += line
            if re.match('^[0-9]-[0-9]', line) != None:
                return [s]


""" Reads a block, which is  one line. Puts the sentence number at the end of the block.
"""
def read_text_block_naacl(stream):
    line = stream.readline()
    pattern = r'^<s snum=([0-9]+)>(.*)</s>$'
    search = re.findall(pattern, line)
    snum = search[0][0]
    string = search[0][1].strip() + ' ' + snum
    return [string]


#### still need to change the output of the following function
#### get rid of sentence number and separate sure and probable alignments


"""
returned string has sentence number at the start, and stream offset at the end.
"""
def read_alignment_block_naacl(stream, offset):
    string = ''
    stream.seek(offset, 0)
    while True:
        line = stream.readline()
        # end of file, append offset info to string and return string
        if not line:
            if string:
                string += str(offset)
                return [string]
            else: return []
        else:
            words = line.strip().split()
            # at the start of a new block
            if string == '':
                snum = words[0]
                string += '-'.join(words[1:]) + ' '
                offset += len(line) # increment offset
            # subsequent lines in a stream block
            else:
                if snum == words[0]:
                    string += '-'.join(words[1:]) + ' '
                    offset += len(line) # increment offset
                # encountering a new block, append offset info
                # to string and return string
                else:
                    string += str(offset)
                    return [string]


# for nltk/corpus/reader/aligned.py

from nltk.corpus.reader.api import CorpusReader
from nltk.corpus.reader.util import StreamBackedCorpusView, concat
from nltk import WhitespaceTokenizer, RegexpTokenizer


class AlignedCorpusReader(CorpusReader):
    """
    This is the base class for giza and naacl readers.
    When using giza or naacl corpus reader, please refer to
    GizaCorpusReader or NaaclCorpusReader class instead.
    """
    def __init__(self, root, fileids, 
                 sep='/', word_tokenizer=WhitespaceTokenizer(),
                 sent_tokenizer=RegexpTokenizer('\n', gaps=True),
                 encoding=None):
        """
        Construct a new Aligned Corpus reader for a set of documents
        located at the given root directory.  Example usage:

            >>> root = '/...path to corpus.../'
            >>> reader = AlignedCorpusReader(root, '.*', '.txt')
        
        @param root: The root directory for this corpus.
        @param fileids: A list or regexp specifying the fileids in this corpus.
        """
        CorpusReader.__init__(self, root, fileids, encoding)
        self._sep = sep
        self._word_tokenizer = word_tokenizer
        self._sent_tokenizer = sent_tokenizer
        self._alignedsent_block_reader=None,
        self._alignedsent_block_reader = self._alignedsent_block_reader
        self._alignedsent_corpus_view = None

    def raw(self, fileids=None):
        """
        @return: the given file(s) as a single string.
        @rtype: C{str}
        """
        if fileids is None: fileids = self._fileids
        elif isinstance(fileids, basestring): fileids = [fileids]
        return concat([self.open(f).read() for f in fileids])

    def words(self, fileids=None):
        """
        @return: the given file(s) as a list of words
            and punctuation symbols.
        @rtype: C{list} of C{str}
        """
        return concat([self._alignedsent_corpus_view(fileid, enc, False, False,
                                             self._word_tokenizer,
                                             self._sent_tokenizer,
                                             self._alignedsent_block_reader)
                       for (fileid, enc) in self.abspaths(fileids, True)])

    def sents(self, fileids=None):
        """
        @return: the given file(s) as a list of
            sentences or utterances, each encoded as a list of word
            strings.
        @rtype: C{list} of (C{list} of C{str})
        """
        return concat([self._alignedsent_corpus_view(fileid, enc, False, True,
                                             self._word_tokenizer,
                                             self._sent_tokenizer,
                                             self._alignedsent_block_reader)
                       for (fileid, enc) in self.abspaths(fileids, True)])

    def aligned_sents(self, fileids=None):
        """
        @return: the given file(s) as a list of AlignedSent objects.
        @rtype: C{list} of C{AlignedSent}
        """
        return concat([self._alignedsent_corpus_view(fileid, enc, True, True,
                                             self._word_tokenizer,
                                             self._sent_tokenizer,
                                             self._alignedsent_block_reader)
                       for (fileid, enc) in self.abspaths(fileids, True)])


class GizaCorpusReader(AlignedCorpusReader):
    """
    Reader for the giza corpus. Tokens are assumed to be separated by
    whitespace.
    """
    def __init__(self, root, fileids, 
                 sep='/', word_tokenizer=WhitespaceTokenizer(),
                 sent_tokenizer=RegexpTokenizer('\n', gaps=True),
                 encoding=None):
        """
        Construct a new Aligned Corpus reader for a set of documents
        located at the given root directory.  Example usage:

            >>> root = '/...path to corpus.../'
            >>> reader = GizaCorpusReader(root, 'alignment-en-fr')
        
        @param root: The root directory for this corpus.
        @param fileids: A list or regexp specifying the fileids in this corpus.
        """
        AlignedCorpusReader.__init__(self, root, fileids, sep, word_tokenizer,
                 sent_tokenizer, encoding='latin-1')
        self._alignedsent_block_reader = read_block_giza
        self._alignedsent_corpus_view = GizaCorpusView


class NaaclCorpusReader(AlignedCorpusReader):
    """
    Reader for the naacl corpus. Readers for text files (trial.e and trial.f)
    and reader for the alignment file (trial.wa) are separated and can be specified in
    initialization with the aligned_filetype argument.

    aligned_filetype can either take 'text' or 'align' as an argument
    """
    def __init__(self, root, fileids, aligned_filetype,
                 sep='/', word_tokenizer=WhitespaceTokenizer(),
                 sent_tokenizer=RegexpTokenizer('\n', gaps=True),
                 encoding='latin-1'):
        """
        Construct a new Aligned Corpus reader for a set of documents
        located at the given root directory.  Example usage:

            >>> root = '/...path to corpus.../'
            >>> reader = NaaclCorpusReader(root, 'trial.wa', 'align')
        
        @param root: The root directory for this corpus.
        @param fileids: A list or regexp specifying the fileids in this corpus.
        """
        AlignedCorpusReader.__init__(self, root, fileids, sep, word_tokenizer,
                 sent_tokenizer, encoding)

        if aligned_filetype == 'text':
            self._alignedsent_block_reader = read_text_block_naacl
            self._alignedsent_corpus_view = NaaclTextCorpusView
        ## filetype == 'align'
        else:
            self._alignedsent_block_reader = read_alignment_block_naacl
            self._alignedsent_corpus_view = NaaclAlignmentCorpusView


from aligned_sent import *

class AlignedCorpusView(StreamBackedCorpusView):

    def __init__(self, corpus_file, encoding, aligned, group_by_sent,
                 word_tokenizer, sent_tokenizer, alignedsent_block_reader):
        self._aligned = aligned
        self._group_by_sent = group_by_sent
        self._word_tokenizer = word_tokenizer
        self._sent_tokenizer = sent_tokenizer
        self._alignedsent_block_reader = alignedsent_block_reader
        StreamBackedCorpusView.__init__(self, corpus_file, encoding=encoding)


class GizaCorpusView(AlignedCorpusView):
    """
    A specialized corpus view for Giza aligned sentences.
    C{AlignedSentCorpusView} objects are typically created by
    L{AlignedCorpusReader} (not directly by nltk users).
    """
        
    def read_block(self, stream):
        block = [self._word_tokenizer.tokenize(sent_str)
                 for alignedsent_str in self._alignedsent_block_reader(stream)
                 for sent_str in self._sent_tokenizer.tokenize(alignedsent_str)]     
        if self._aligned:
            block[2] = " ".join(block[2]) # kludge; we shouldn't have tokenized the alignment string
            block = [AlignedSent(*block)]
        elif self._group_by_sent:
            block = [block[0]]
        else:
            block = block[0]

        return block


"""
This Corpus View is for handling naacl alignment file, for example trial.wa
"""
class NaaclAlignmentCorpusView(AlignedCorpusView):
    """
    A specialized corpus view for Naacl alignment.
    C{AlignedSentCorpusView} objects are typically created by
    L{AlignedCorpusReader} (not directly by nltk users).
    """
    def __init__(self, corpus_file, encoding, aligned, group_by_sent,
                 word_tokenizer, sent_tokenizer, alignedsent_block_reader):
        AlignedCorpusView.__init__(self, corpus_file, encoding, aligned, group_by_sent,
                 word_tokenizer, sent_tokenizer, alignedsent_block_reader)
        self._infile_offset = 0
        
    def read_block(self, stream):
        sent = self._alignedsent_block_reader(stream, self._infile_offset)
        block = [self._word_tokenizer.tokenize(sent_str)
                 for sent_str in sent]
        # get the current offset
        self._infile_offset = int(block[0].pop())
        if self._group_by_sent:
            block = [block[0]]
        else:
            block = block[0]
        return block


"""
This Corpus View is for handling naacl text files, for example trial.e and trial.f
"""
class NaaclTextCorpusView(AlignedCorpusView):
    """
    A specialized corpus view for Naacl aligned sentences.
    C{AlignedSentCorpusView} objects are typically created by
    L{AlignedCorpusReader} (not directly by nltk users).
    """
        
    def read_block(self, stream):
        sent = self._alignedsent_block_reader(stream)
        block = [self._word_tokenizer.tokenize(sent_str)
                 for sent_str in sent]
        if self._group_by_sent:
            block = [block[0]]
        else:
            block = block[0]
        return block

